from pathlib import Path
import json
from functools import partial
from concurrent.futures import ThreadPoolExecutor, Future
from typing import Any, Tuple

import numpy as np
import cv2
from superstream import Stream

from .common import image_suffix


font_face = cv2.FONT_HERSHEY_COMPLEX_SMALL
font_scale = 1


def _json2yolo(json_path: Path, output_path: Path, labels_info: dict) -> None:
    with json_path.open() as fp:
        data = json.load(fp)

    def normalize(data_shapes: dict, w, h):
        points = data_shapes["points"]
        x1, y1 = points[0]
        x2, y2 = points[1]
        x, y = (x1 + x2) / 2 / w, (y1 + y2) / 2 / h
        w, h = abs(x2 - x1) / w, abs(y2 - y1) / h

        label_idx = labels_info.setdefault(data_shapes["label"], len(labels_info))
        return f"{label_idx} {x:.4f} {y:.4f} {w:.4f} {h:.4f}"

    w, h = data["imageWidth"], data["imageHeight"]
    cur_normalize = partial(normalize, w=w, h=h)
    output_path.write_text(
        "\n".join(
            map(
                cur_normalize,
                filter(lambda data_dict: data_dict["shape_type"] == "rectangle", data["shapes"]),
            )
        )
    )


def json2yolo(json_folder: Path, output_folder: Path):
    output_folder.mkdir(exist_ok=True)
    labels_info = {}
    for json_file in json_folder.glob("*.json"):
        output_file = output_folder / (json_file.stem + ".txt")
        _json2yolo(json_file, output_file, labels_info)

    (output_folder.resolve().parent / "label_info.txt").write_text(
        "\n".join(
            Stream(labels_info.items()).map(lambda itms: f"{itms[0]}\t=\t{itms[1]}").to_list()
        )
    )


def _imwrite(path, img):
    cv2.imencode(".png", img)[1].tofile(path)


# 坐标转换，原始存储的是YOLOv5格式
# Convert nx4 boxes from [x, y, w, h] normalized to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
def _xywh2xyxy(yolo_content, img_w, img_h, img, output_path):
    label_idx, x, y, w, h = yolo_content
    label = label_idx
    # 边界框反归一化
    x_t = float(x) * img_w
    y_t = float(y) * img_h
    w_t = float(w) * img_w
    h_t = float(h) * img_h

    # 计算坐标
    top_left_x = x_t - w_t / 2
    top_left_y = y_t - h_t / 2
    bottom_right_x = x_t + w_t / 2
    bottom_right_y = y_t + h_t / 2

    # 绘图  rectangle()函数需要坐标为整数
    pt1 = (int(top_left_x), int(top_left_y))
    pt2 = (int(bottom_right_x), int(bottom_right_y))
    cv2.rectangle(
        img,
        pt1,
        pt2,
        (71, 127, 255),
        2,
    )
    _, base_line = cv2.getTextSize(label, font_face, font_scale, 1)
    cv2.putText(img, label, (pt1[0], pt1[1] - base_line), font_scale, font_face, (0, 0, 0), 1)

    _imwrite(str(output_path), img)


def _inner_process(result: Future[tuple[Any, Path, Path, Path]]):
    img, img_path, lable_dir, output_dir = result.result()
    print(img_path)
    h, w = img.shape[:2]

    img_stem = img_path.stem
    label_path = lable_dir / (img_stem + ".txt")
    img_name = img_path.name
    output_path = output_dir / img_name

    with open(label_path) as fp:
        (
            Stream(fp.readlines())
            .map(lambda txt: txt.strip().split())
            .for_each(lambda txt: _xywh2xyxy(txt, w, h, img, output_path))
        )


def _imread(img_path: str, cv_flag):
    """用于解决opencv库中imread无法读取含有中文路径的imread函数"""
    return cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv_flag)


def visual_by_yolo(img_dir: Path, lable_dir: Path, output_dir: Path):
    with ThreadPoolExecutor() as exe:
        (
            Stream(img_dir.glob("*"))
            .map(
                lambda i_p: exe.submit(
                    lambda p: (_imread(str(p), cv2.IMREAD_ANYCOLOR), p, lable_dir, output_dir),
                    i_p,
                )
            )
            .for_each(lambda f: f.add_done_callback(_inner_process))
        )


def _inner_resize(result: Future[tuple[Any, Path, Path, Tuple[float, float]]]):
    img, img_path, dst_dir, new_sc = result.result()
    scale_x, scale_y = new_sc

    _imwrite(dst_dir / img_path.name, cv2.resize(img, None, fx=scale_x, fy=scale_y))


def resize_img(src: Path, dst: Path, new_scale: Tuple[float, float]):
    with ThreadPoolExecutor() as exe:
        (
            Stream(src.glob("*"))
            .filter(lambda sfx: sfx in image_suffix)
            .map(
                lambda i_p: exe.submit(
                    lambda p: (_imread(str(p), cv2.IMREAD_ANYCOLOR), p, dst, new_scale), i_p
                )
            )
            .for_each(lambda f: f.add_done_callback(_inner_resize))
        )


"""
AI生成的代码
"""

import os
import shutil
import random

pic_suffix = image_suffix


def split_data(
    image_folder: str,
    label_folder: str,
    train_ratio=0.7,
    test_ratio=0.2,
    output_dir=".",
):
    # 创建存储图片和标签的文件夹
    os.chdir(output_dir)
    os.makedirs("images", exist_ok=True)
    os.makedirs("labels", exist_ok=True)
    os.makedirs("images/train", exist_ok=True)
    os.makedirs("images/test", exist_ok=True)
    os.makedirs("images/val", exist_ok=True)
    os.makedirs("labels/train", exist_ok=True)
    os.makedirs("labels/test", exist_ok=True)
    os.makedirs("labels/val", exist_ok=True)

    # 获取图片文件名列表
    image_files = [f for f in os.listdir(image_folder) if any(f.endswith(s) for s in pic_suffix)]
    label_files = [f[:-4] + ".txt" for f in image_files]

    # assert len(image_files) == len(label_files), "图片和标签文件数量不匹配"
    # assert all(i[0][:-4] == i[1][:-4] for i in zip(image_files, label_files)), "图片与标签不对应"

    # 随机打乱图片文件名列表
    idx = list(range(len(image_files)))
    random.shuffle(idx)
    image_files = [image_files[i] for i in idx]
    label_files = [label_files[i] for i in idx]

    # 计算每个集合的大小
    total_size = len(image_files)
    train_size = int(total_size * train_ratio)
    test_size = int(total_size * test_ratio)

    # 将图片和标签文件分配到各个集合中
    train_images = image_files[:train_size]
    train_labels = label_files[:train_size]
    test_images = image_files[train_size : train_size + test_size]
    test_labels = label_files[train_size : train_size + test_size]
    val_images = image_files[train_size + test_size :]
    val_labels = label_files[train_size + test_size :]

    # 将图片和标签文件复制到相应的文件夹中
    for img, lbl in zip(train_images, train_labels):
        shutil.copy(os.path.join(image_folder, img), "images/train")
        shutil.copy(os.path.join(label_folder, lbl), "labels/train")
    for img, lbl in zip(test_images, test_labels):
        shutil.copy(os.path.join(image_folder, img), "images/test")
        shutil.copy(os.path.join(label_folder, lbl), "labels/test")
    for img, lbl in zip(val_images, val_labels):
        shutil.copy(os.path.join(image_folder, img), "images/val")
        shutil.copy(os.path.join(label_folder, lbl), "labels/val")


# 调用函数进行数据拆分
# split_data(r"担架", r"担架_yolo")
